import base64
import time
import hashlib

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


def sha256_hash(data):
    sha256 = hashlib.sha256()
    sha256.update(data)
    return sha256.digest()


def generate_open_id_and_second_key(private_key_pem, partner_id):
    private_key = serialization.load_pem_private_key(
        private_key_pem.encode('utf-8'),
        password=None,
        backend=default_backend()
    )
    private_key_p = private_key.private_numbers().p.to_bytes(128, 'big')    # 私钥大质数p、q
    aes_key = private_key_p[0:32]       # 256位（32字节）AES密钥
    aes_iv = private_key_p[32:48]       # 生成随机的初始化向量
    hash_tail = private_key_p[48:112]   # 生成second_key用64字节随机数尾
    timestamp = int(time.time() * 1000000).to_bytes(8, 'big')
    open_id = Cipher(
        algorithms.AES(aes_key),
        modes.CBC(aes_iv),
        backend=default_backend()
    ).encryptor().update(partner_id.to_bytes(8, 'big') + timestamp)
    second_key = sha256_hash(open_id + hash_tail)
    open_id = base64.b64encode(open_id).decode('utf-8')
    second_key = base64.b64encode(second_key).decode('utf-8')
    return open_id, second_key


def generate_private_and_public_key():
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    public_key = private_key.public_key()
    private_key_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    ).decode('utf-8')
    public_key_pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    ).decode('utf-8')
    return private_key_pem, public_key_pem


def sign(private_key_pem, message):
    private_key = serialization.load_pem_private_key(
        private_key_pem.encode('utf-8'),
        password=None,
        backend=default_backend()
    )
    signature = private_key.sign(
        data=message.encode('utf-8'),
        padding=padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        algorithm=hashes.SHA256()
    )
    return base64.b64encode(signature).decode('utf-8')


def verify_signature(public_key_pem, signature, data):
    signature_bytes = base64.b64decode(signature.encode('utf-8'))
    public_key = serialization.load_pem_public_key(
        public_key_pem.encode('utf-8'),
        backend=default_backend()
    )
    public_key.verify(
        signature=signature_bytes,
        data=data.encode('utf-8'),
        padding=padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        algorithm=hashes.SHA256()
    )
